#include "Riostream.h"
#include <iostream>
#include <iomanip>
#include <stdexcept>
#include <string>

#include "TFile.h"
#include "TTree.h"
#include "TChain.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TCanvas.h"
#include "TPad.h"
#include "TStyle.h"
#include "TROOT.h"
#include "TMath.h"
#include "TRandom.h"
#include "TFormula.h"
#include "TSystem.h"
#include "TEnv.h"

#include "CommandLineInterface.hh"
#include "Barrel.hh"
#include "Annular.hh"
#include "Germanium.hh"
#include "SRIMloader.hh"

using std::cerr;
using std::cout;
using std::endl;

ClassImp(Barrel);
ClassImp(Annular);
ClassImp(Germanium);

SRIM srim;

int const kCoincidenceTime = 1000;

double g_angles[8][3][6][2];
int64_t dt;
float ge_e;
float ge_dc_e;
int64_t ge_t;
int16_t ge_id;

void FillWithCoincidentGammas(TTree *tree, vector<Germanium> *Miniball, double
    time, double p_x, double p_y, double p_z)
{
  // Gammas.
  for (size_t j = 0; Miniball->size() > j; ++j) {
    for (size_t k = 0; (*Miniball)[j].GetCrystal().size() > k; ++k) {
      ge_t = (*Miniball)[j].GetCrystal()[k].GetTime() * 25;
      dt = time - ge_t;
      if (kCoincidenceTime > abs(dt)) {

        double max_E = 0.0;
        int max_seg = 0;
        for (size_t l = 0; (*Miniball)[j].GetCrystal()[k].GetSeg().size() > l;
            ++l) {
          if (max_E < (*Miniball)[j].GetCrystal()[k].GetSeg()[l]) {
            max_E = (*Miniball)[j].GetCrystal()[k].GetSeg()[l];
            max_seg = (*Miniball)[j].GetCrystal()[k].GetSegID()[l];
          }
        }

        int clu_id = (*Miniball)[j].GetCluID();
        int cry_id = (*Miniball)[j].GetCrystal()[k].GetCryID();
        int seg_id = max_seg;

        double g_theta = g_angles[clu_id][cry_id][seg_id][0];
        double g_phi = g_angles[clu_id][cry_id][seg_id][1];

        double g_x = sin(g_phi) * sin(g_theta);
        double g_y = cos(g_phi) * sin(g_theta);
        double g_z = cos(g_theta);

        double cos_theta = p_x * g_x + p_y * g_y + p_z * g_z;

        double gamma = 1;
        double beta = 0;

        ge_e = (*Miniball)[j].GetCrystal()[k].GetCore();
        ge_dc_e = ge_e * gamma * (1 - beta * cos_theta);
        ge_id = 100 * clu_id + 10 * cry_id + seg_id;
        tree->Fill();
      }
    }
  }
}

void LoadAngles(std::string const &file_name)
{
  std::ifstream file(file_name.c_str());
  if (!file.is_open()) {
    throw std::runtime_error("Could not load angles file.");
  }
  for (int clu_id = 0; 8 > clu_id; ++clu_id) {
    for (int cry_id = 0; 3 > cry_id; ++cry_id) {
      for (int seg_id = 0; 6 > seg_id; ++seg_id) {
        for (int coord = 0; 2 > coord; ++coord) {
          if (!file.good()) {
            throw std::runtime_error("Error parsing angles file.");
          }
          file >> g_angles[clu_id][cry_id][seg_id][coord];
        }
      }
    }
  }
}

int main(int argc, char **argv)
{
  vector<char*> InputFiles;
  char *OutputFile = NULL;
  char *MBFile = NULL;
  char *cutfile = NULL;
  char *srimdir = NULL;
  bool verbose = false;

  float ttarget, thetain;
  float Rafter, Ein;
  float Eex[3][5], gamma[3][5], beta[3][5];
  float m[30];
  int par[3][5];
  int Np, Nd, Nt, Np_FB, Nd_FB, Nt_FB, Np_CD, Nd_CD, Nt_CD;
  float time[3][5];
  float CDThreshold[4];

  // Masses
  m[1] = 1.00782503207*931.494e3;
  m[2] = 2.01410177785*931.494e3;
  m[3] = 3.01604927767*931.494e3;
  m[4] = 4.00260325415*931.494e3;
  m[10] = 10.013533818*931.494e3;
  m[11] = 11.021657749*931.494e3;
  m[12] = 12.026920737*931.494e3;
  m[21] = 20.993846684*931.494e3;
  m[22] = 21.991385113*931.494e3;  // 22Ne
  //m[22] = 22.002998815*931.494e3;  // 22F
  m[23] = 22.994466904*931.494e3;

  // Beam energy
  Ein = 2.85e3*11;
  thetain = 0.;
  ttarget = 10.;

  // Program argument parsing.
  CommandLineInterface* interface = new CommandLineInterface();
  interface->Add("-i", "inputfiles", &InputFiles);
  interface->Add("-o", "outputfile", &OutputFile);
  interface->Add("-M", "miniball angle", &MBFile);
  interface->Add("-C", "cut functions", &cutfile);
  interface->Add("-S", "SRIM files directory", &srimdir);
  interface->Add("-v", "verbose", &verbose);
  interface->CheckFlags(argc, argv);
  if (0 == InputFiles.size() || 0 == OutputFile) {
    cerr << "You have to provide at least one input file and the output file!"
        << endl;
    return 1;
  }
  if (0 == srimdir) {
    cerr << "SRIM directory is needed!" << endl;
    return 1;
  }

  cout << "input file(s):" << endl;
  for (size_t i = 0;  i < InputFiles.size();  ++i) {
    cout << InputFiles[i] << endl;
  }
  cout << "output file: " << OutputFile << endl;

  // Create SRIM graphs for energy loss calculations.
  srim.Creategraphs(srimdir);
  //printf("R = %f, E = %f \n",srim.RMylar[4]->Eval(3.0),srim.EMylar[4]->Eval(11.0));
  Rafter = srim.RT[11]->Eval(Ein);
  Ein = srim.ET[11]->Eval(Rafter - ttarget/(2.*cos(thetain)));

  LoadAngles(MBFile);

  TChain *tr = new TChain("caltr");
  for (size_t i = 0;  i < InputFiles.size();  ++i) {
    tr->Add(InputFiles[i]);
  }

  if (0 == tr) {
    cout << "could not find tree caltr in file." << endl;
    for (size_t i = 0;  i < InputFiles.size();  ++i) {
      cout << InputFiles[i] << endl;
    }
    return 3;
  }

  vector<Barrel> *FBarrel = new vector<Barrel>;
  vector<Barrel> *BBarrel = new vector<Barrel>;
  vector<Annular> *FCD = new vector<Annular>;
  vector<Germanium> *Miniball = new vector<Germanium>;
  long long EbisTime;
  long long T1Time;
  long long SuperCycleTime;

  tr->SetBranchAddress("ForwardBarrel",&FBarrel);
  tr->SetBranchAddress("BackwardBarrel",&BBarrel);
  tr->SetBranchAddress("ForwardCD",&FCD);
  tr->SetBranchAddress("Miniball",&Miniball);
  tr->SetBranchAddress("EbisTime",&EbisTime);
  tr->SetBranchAddress("T1Time",&T1Time);
  tr->SetBranchAddress("SuperCycleTime",&SuperCycleTime);

  TFile* outfile = new TFile(OutputFile,"recreate");
  if (outfile->IsZombie()) {
    return 4;
  }
  TTree *tree = new TTree("is430", "is430");

  int64_t ev_id;
  int64_t ebis_t;
  int64_t t1_t;
  tree->Branch("id", &ev_id, "id/L");
  tree->Branch("ebis_t", &ebis_t, "ebis_t/L");
  tree->Branch("t1_t", &t1_t, "t1_t/L");
  tree->Branch("dt", &dt, "dt/L");

  tree->Branch("ge_e", &ge_e, "ge_e/F");
  tree->Branch("ge_dc_e", &ge_dc_e, "ge_dc_e/F");
  tree->Branch("ge_t", &ge_t, "ge_t/I");
  tree->Branch("ge_id", &ge_id, "ge_id/s");

  float fb_de_raw;
  float fb_de;
  float fb_e_raw;
  float fb_e;
  float fb_res_raw;
  float fb_res;
  int64_t fb_t;
  int8_t fb_det_id;
  int8_t fb_strip_id;
  float fb_strip_pos;
  float fb_theta;
  float fb_phi;
  float fb_p_de;
  float fb_p_e;
  tree->Branch("fb_de_raw", &fb_de_raw, "fb_de_raw/F");
  tree->Branch("fb_de", &fb_de, "fb_de/F");
  tree->Branch("fb_e_raw", &fb_e_raw, "fb_e_raw/F");
  tree->Branch("fb_e", &fb_e, "fb_e/F");
  tree->Branch("fb_res_raw", &fb_res_raw, "fb_res_raw/F");
  tree->Branch("fb_res", &fb_res, "fb_res/F");
  tree->Branch("fb_t", &fb_t, "fb_t/I");
  tree->Branch("fb_det_id", &fb_det_id, "fb_det_id/b");
  tree->Branch("fb_strip_id", &fb_strip_id, "fb_strip_id/b");
  tree->Branch("fb_strip_pos", &fb_strip_pos, "fb_strip_pos/F");
  tree->Branch("fb_theta", &fb_theta, "fb_theta/F");
  tree->Branch("fb_phi", &fb_phi, "fb_phi/F");
  tree->Branch("fb_p_de", &fb_p_de, "fb_p_de/F");
  tree->Branch("fb_p_e", &fb_p_e, "fb_p_e/F");

  float bb_de;
  float bb_e;
  int64_t bb_t;
  int8_t bb_det_id;
  int8_t bb_strip_id;
  float bb_strip_pos;
  float bb_theta;
  float bb_phi;
  float bb_p_de;
  float bb_p_e;
  tree->Branch("bb_de", &bb_de, "bb_de/F");
  tree->Branch("bb_e", &bb_e, "bb_e/F");
  tree->Branch("bb_t", &bb_t, "bb_t/I");
  tree->Branch("bb_det_id", &bb_det_id, "bb_det_id/b");
  tree->Branch("bb_strip_id", &bb_strip_id, "bb_strip_id/b");
  tree->Branch("bb_strip_pos", &bb_strip_pos, "bb_strip_pos/F");
  tree->Branch("bb_theta", &bb_theta, "bb_theta/F");
  tree->Branch("bb_phi", &bb_phi, "bb_phi/F");
  tree->Branch("bb_p_de", &bb_p_de, "bb_p_de/F");
  tree->Branch("bb_p_e", &bb_p_e, "bb_p_e/F");

  float fcd_e;
  float fcd_p_e;
  int64_t fcd_t;
  int8_t fcd_det_id;
  float fcd_ring_de;
  float fcd_ring_p_de;
  int8_t fcd_ring_id;
  float fcd_strip_de;
  float fcd_strip_p_de;
  int8_t fcd_strip_id;
  float fcd_theta;
  float fcd_phi;
  tree->Branch("fcd_e", &fcd_e, "fcd_e/F");
  tree->Branch("fcd_p_e", &fcd_p_e, "fcd_p_e/F");
  tree->Branch("fcd_t", &fcd_t, "fcd_t/I");
  tree->Branch("fcd_det_id", &fcd_det_id, "fcd_det_id/b");
  tree->Branch("fcd_ring_de", &fcd_ring_de, "fcd_ring_de/F");
  tree->Branch("fcd_ring_p_de", &fcd_ring_p_de, "fcd_ring_p_de/F");
  tree->Branch("fcd_ring_id", &fcd_ring_id, "fcd_ring_id/b");
  tree->Branch("fcd_strip_de", &fcd_strip_de, "fcd_strip_de/F");
  tree->Branch("fcd_strip_p_de", &fcd_strip_p_de, "fcd_strip_p_de/F");
  tree->Branch("fcd_strip_id", &fcd_strip_id, "fcd_strip_id/b");
  tree->Branch("fcd_theta", &fcd_theta, "fcd_theta/F");
  tree->Branch("fcd_phi", &fcd_phi, "fcd_phi/F");

  CDThreshold[0] = 500.;
  CDThreshold[1] = 1300.;
  CDThreshold[2] = 500.;
  CDThreshold[3] = 800.;

  // Create cut functions.
  TFormula* For_cut[4][16];
  TFormula* Back_cut[4];
  TFormula* CD_cut[6][4];
  TFormula* CD_PT_cut[6][4];

  TEnv *Cutfunc = new TEnv(cutfile);

  // Forward barrel.
  for (int istrip = 0; istrip < 16; istrip++) {
    for (int icut = 0; icut < 4; icut++) {
      For_cut[icut][istrip] = new TFormula(
          Form("For_cut_%d_%d",icut,istrip),
          Cutfunc->GetValue(Form("For.%d.%d",icut,istrip), "0*x"));
    }
  }

  // Backward barrel.
  for (int icut = 0; icut < 4; icut++) {
    Back_cut[icut] = new TFormula(
        Form("Back_cut_%d",icut),
        Cutfunc->GetValue(Form("Back.%d",icut), "0*x"));
  }

  // CD.
  for (int idet = 0; idet < 4; idet++) {
    for (int icut = 0; icut < 6; icut++) {
      CD_cut[icut][idet] = new TFormula(
          Form("CD_cut_%d_%d",icut,idet),
          Cutfunc->GetValue(Form("CD.%d.%d",icut,idet), "0*x"));
      CD_PT_cut[icut][idet] = new TFormula(
          Form("CD_PT_cut_%d_%d",icut,idet),
          Cutfunc->GetValue(Form("CD_PT.%d.%d",icut,idet), "0*x"));
    }
  }

  int64_t nentries = tr->GetEntries();
  int nbytes = 0;
  int status;

  Nd = 0;
  Np = 0;
  Nt = 0;
  Nd_FB = 0;
  Np_FB = 0;
  Nt_FB = 0;
  Nd_CD = 0;
  Np_CD = 0;
  Nt_CD = 0;

  for (int64_t i = 0; nentries > i; ++i) {
    for (int idet = 0; 3 > idet; idet++) {
      for (int ie = 0; 5 > ie; ie++) {
        par[idet][ie] = 0;
        time[idet][ie] = 0;
        Eex[idet][ie] = 0;
        gamma[idet][ie] = 0;
        beta[idet][ie] = 0;
      }
    }
    status = tr->GetEvent(i);
    if (-1 == status) {
      cerr << "Error occured, couldn't read entry " << i << " from tree " <<
          tr->GetName() << " in file " << tr->GetFile()->GetName() << endl;
      return 5;
    }
    else if (0 == status) {
      cerr << "Error occured, entry " << i << " in tree " << tr->GetName() <<
          " in file " << tr->GetFile()->GetName() << " doesn't exist" << endl;
      return 6;
    }
    nbytes += status;

    ev_id = i;
    ebis_t = EbisTime * 25;
    t1_t = T1Time * 25;

    bb_de = -1;
    fcd_e = -1;

    // Multiplicities ok!

    // Forward barrel.
    //  Find the strip with the highest signal.
    //  Record strip and the energy.
    //  Should check cross talk in MyCal_histos.cc.
    for (size_t j = 0; FBarrel->size() > j; ++j) {
      if (0 == (*FBarrel)[j].GetMult()) {
        continue;
      }
      fb_de = (*FBarrel)[j].GetRear();
      fb_e = (*FBarrel)[j].GetEdet();
//      fb_res = (*Barrel)[j].GetStrip();
      fb_t = (*FBarrel)[j].GetTime() * 25;

      fb_det_id = (*FBarrel)[j].GetID();
      fb_strip_id = (*FBarrel)[j].GetStripNr()[0];
      fb_strip_pos = (*FBarrel)[j].GetStripPos()[0];

      // Proton angle.
      double z = 8.0 + 3.125 * (fb_strip_id + gRandom->Rndm());
      double y = 29.0;
      double x = 50.0 * (fb_strip_pos - 0.5);
      double f = 1 / sqrt(x * x + y * y + z * z);
      x *= f;
      y *= f;
      z *= f;
      fb_theta = acos(z);
      fb_phi = atan(x / y) + TMath::Pi() * 0.5 * (fb_det_id + 1);

      fb_p_de = fb_de * sin(fb_theta);
      fb_p_e = fb_e + (fb_de - fb_p_de);

      FillWithCoincidentGammas(tree, Miniball, fb_t, x, y, z);
    }

    fb_de = -1;

    // Backward barrel.
    for (size_t j = 0; BBarrel->size() > j; ++j) {
      if (0 == (*BBarrel)[j].GetStripNr().size()) {
        continue;
      }
      bb_t = (*BBarrel)[j].GetTime() * 25;
      bb_de = (*BBarrel)[j].GetRear();
      bb_e = (*BBarrel)[j].GetEdet();

      bb_det_id = (*BBarrel)[j].GetID();
      bb_strip_id = (*BBarrel)[j].GetStripNr()[0];
      bb_strip_pos = (*BBarrel)[j].GetStripPos()[0];

      double z = -(8. + 3.125 * (bb_strip_id + gRandom->Rndm()));
      double y = 29.;
      double x = 50. * (bb_strip_pos - .5);
      bb_theta = acos(z / sqrt(x * x + y * y + z * z));
      bb_phi = atan(x / y) + TMath::Pi()/2. * (bb_det_id + 1);

      bb_p_de = bb_de * sin(bb_theta);
      bb_p_e = bb_e + (bb_de - bb_p_de);

      FillWithCoincidentGammas(tree, Miniball, bb_t, x, y, z);
    }

    bb_de = -1;

    // Forward CD.
    for (size_t j = 0; FCD->size() > j; ++j) {
      fcd_t = (*FCD)[j].GetTime() * 25;
      fcd_e = (*FCD)[j].GetEdet();
      fcd_det_id = (*FCD)[j].GetID();

      fcd_ring_de = -1;
      for (size_t k = 0; (*FCD)[j].GetRingNr().size() > k; ++k) {
        if (fcd_ring_de < (*FCD)[j].GetRingEnergy()[k]) {
          fcd_ring_id = (*FCD)[j].GetRingNr()[k];
          fcd_ring_de = (*FCD)[j].GetRingEnergy()[k];
        }
      }

      fcd_strip_de = -1;
      for (size_t k = 0; (*FCD)[j].GetStripNr().size() > k; ++k) {
        if (fcd_strip_de < (*FCD)[j].GetStripEnergy()[k]) {
          fcd_strip_id = (*FCD)[j].GetStripNr()[k];
          fcd_strip_de = (*FCD)[j].GetStripEnergy()[k];
        }
      }

      fcd_theta = atan((9.0 + (gRandom->Rndm() + fcd_ring_id) * 2.0) / 63.0);
      if (4 <= fcd_strip_id && 12 > fcd_strip_id) {
        fcd_phi = .0593411 * (4 + 2 * (fcd_strip_id - 4 + gRandom->Rndm())) - .71209 + fcd_det_id * TMath::Pi() / 2.;
      }
      else {
        fcd_phi = .0593411 * (fcd_strip_id + gRandom->Rndm()) - .71209 + fcd_det_id * TMath::Pi() / 2.;
      }

      double cos_fcd_theta = cos(fcd_theta);
      fcd_ring_p_de = fcd_ring_de * cos_fcd_theta;
      fcd_strip_p_de = fcd_strip_de * cos_fcd_theta;
      fcd_p_e = fcd_e + (fcd_ring_de - fcd_ring_p_de) + (fcd_strip_de - fcd_strip_p_de);

//      FillWithCoincidentGammas(tree, Miniball, fcd_t);
    }

    if (0 == i % 1000) {
      cout << setw(5) << setiosflags(ios::fixed) << setprecision(1) << (100.*i)/nentries << " % done\r" << flush;
    }
  }

  outfile->Write();
  outfile->Close();

  return 0;
}
